use std::iter::zip;

use candle_core::Tensor;

mod embd;
pub use embd::Embd;

pub use candle_core::{DType, Device};

/// 测量一维向量`a`与`b`之间的欧式距离。并不是好的实现。
pub fn dis(a: &Tensor, b: &Tensor) -> candle_core::Result<f32> {
    let dim_a: Vec<f32> = a.to_vec1()?;
    let dim_b: Vec<f32> = b.to_vec1()?;
    assert_eq!(dim_a.len(), dim_b.len());
    let mut dis = 0f32;
    for (da, db) in zip(dim_a, dim_b) {
        dis += (da - db).powi(2);
    }
    Ok(dis)
}

#[cfg(test)]
mod tests {
    use super::*;
    use anyhow::Result;
    use candle_core::IndexOp;
    use embedding_model::EmbeddingModel;

    #[test]
    fn test_load() -> Result<()> {
        let _ = Embd::load(
            r"model\text2vec-base-chinese_safetensors",
            candle_core::DType::F32,
            &candle_core::Device::Cpu,
        )?;
        Ok(())
    }

    #[test]
    fn test_get_single_embedding() -> Result<()> {
        let embd = Embd::load(
            r"model\text2vec-base-chinese_safetensors",
            candle_core::DType::F32,
            &candle_core::Device::Cpu,
        )?;
        let res = embd.embedding_single_tensor("hello")?;
        println!("res = {}", res);
        Ok(())
    }

    #[test]
    fn test_get_batch_embedding() -> Result<()> {
        let embd = Embd::load(
            r"model\text2vec-base-chinese_safetensors",
            candle_core::DType::F32,
            &candle_core::Device::Cpu,
        )?;
        let sentences = ["hello", "你好","hello","hi"];
        let res = embd.embedding_batch_tensor(sentences.to_vec())?;
        println!("res = {}", res);
        println!("dis 2 1 = {}", dis(&res.i(0)?, &res.i(1)?)?);
        println!("dis 3 1 = {}", dis(&res.i(0)?, &res.i(2)?)?);
        println!("dis 3 2 = {}", dis(&res.i(1)?, &res.i(2)?)?);
        println!("dis 3 4 = {}", dis(&res.i(3)?, &res.i(2)?)?);
        Ok(())
    }

    #[test]
    fn test_get_batch_embedding_vec() -> Result<()> {
        let embd = Embd::load(
            r"model\text2vec-base-chinese_safetensors",
            candle_core::DType::F32,
            &candle_core::Device::Cpu,
        )?;
        let sentences = ["hello", "你好"];
        let res = embd.embedding_batch_vec(sentences.to_vec())?;
        println!("{:?}", res);
        Ok(())
    }
}
